Chat-template repair: warn-by-default, AST classification, dict support#49
Chat-template repair: warn-by-default, AST classification, dict support#49danielhanchen wants to merge 13 commits into
Conversation
Follow-up hardening on top of PR unslothai#4426 (which fixed the unslothai#4150 RuntimeError for ChatML LoRA reloads). Behavior changes: - Warn-by-default instead of RuntimeError. When fix_chat_template cannot repair a broken template, emit a warning and return the original. Set UNSLOTH_STRICT_CHAT_TEMPLATE=1 to restore the pre-warn hard fail. Fixes the UX where a missing `{% if add_generation_prompt %}` block on a saved LoRA (typical after LlamaFactory / Axolotl re-serialize) would block model loading entirely. - Local path vs HF hub distinguished in the warning message. For local paths the message points at the likely downstream tool; for HF IDs it points at the upstream model maintainers. Previously both said "file a bug report to the maintainers of <path>" even when <path> was the user's own saves/ directory. - Dict / list chat_template now handled. Hermes-3 ships with {default, tool_use} and the previous code crashed with AttributeError: 'dict' object has no attribute 'find' when entering _fix_chat_template with a dict. Each variant is now fixed independently; structure is preserved. Internals: - _find_end_position now matches all four Jinja whitespace-control variants ({% %}, {%- %}, {% -%}, {%- -%}) and returns the rightmost endfor/endif so multi-for templates aren't locked onto the first loop. Previously {%- endfor -%} (both-side dash, used by Qwen3-Guard) was silently bypassed. - _has_add_generation_prompt_block uses Jinja AST via jinja2.nodes.If/Name walks instead of substring matching, so templates that hide the block behind comments or dash-style variants are classified correctly. - _template_ends_with_toplevel_for gates the GH#4150 ChatML repair on the AST: only fires when the last structural top-level node is a For (standard ChatML shape), ignoring trailing pure-whitespace output nodes. Templates wrapped in an outer If (Qwen3-Guard) are now explicitly skipped at the _fix_chat_template level as well, not just at load_correct_tokenizer's name-based exemption. - _validate_patched_template renders the patched template with and without add_generation_prompt and confirms the patched output responds to the flag by appending (not replacing) content. If validation fails, the patch is discarded and we fall through to the warn path. Verified with an expanded regression suite in tests/: - test_fix_chat_template_pr4426.py: 42/42 template-matrix cells - test_load_correct_tokenizer_pr4426.py: 5/5 tokenizer loads - test_chat_template_followups.py: 10/10 new follow-up tests - test_mistral_pr4426.py: 5 Mistral variants byte-identical - test_qwen_pr4426.py: 14 Qwen variants byte-identical (Qwen1.5, Qwen2, Qwen2.5-Instruct/Coder/Math/VL, Qwen3, Qwen3-Coder, QwQ, Qwen3-Guard-Gen)
for more information, see https://pre-commit.ci
If tokenizer.chat_template is a property or otherwise read-only, the validation helper would crash with AttributeError when trying to temporarily set the patched template. Catch the assignment failure and return False (skip validation), and best-effort restore in the finally block.
… non-ChatML templates
The previous `_infer_assistant_separator` was a four-tier regex heuristic that
only worked on ChatML-shaped templates and forced a hard `<|im_start|>` /
`<|im_end|>` presence gate on Case 2 repair. This meant a Llama-3, Gemma, or
Phi-3 template stripped of its generation-prompt block by a downstream tool
(LlamaFactory, Axolotl, etc.) would still warn-and-return even though the
structural shape is identical to the ChatML case the PR already handles.
This replaces the regex with `_derive_assistant_prefix_by_render`: render the
template with two dialogs that differ only in assistant content, then
`os.path.commonprefix` on the tails captures the exact assistant-turn prefix
the template emits. The template itself is ground truth, so non-ChatML shapes
work as long as the assistant block is a literal the template emits once per
message.
Three guards keep the derivation safe:
A. both assistant renders extend the base render (no reordering);
B. the divergence point is exactly the content-insertion site (sentinel
follows the common prefix);
C. a user-role cross-check: if a render with a user sentinel also emits
the same prefix, role has no effect on output and we reject. A render
failure on [user, user] (e.g. Gemma's `raise_exception` alternation
check) is evidence that role matters; we accept.
Sentinels differ at character 0 so `commonprefix` cannot absorb them, and
trailing whitespace/comments after the last `{% endfor %}` are stripped
before probing (they would appear in base but not after the appended
assistant turn and break Guard A).
`_fix_chat_template` and `_repair_string_template` now thread an
`is_sharegpt` kwarg; `_fix_chat_template` retries once with
`is_sharegpt=True` if the first probe returns None (dual-probe fallback
for dict/list callers).
The ChatML `<|im_start|>` / `<|im_end|>` hard gate in Case 2 is dropped.
`_infer_assistant_separator` is deleted.
Verified via:
- tests/test_fix_chat_template_pr4426.py: 51/51 cells (new Llama-3,
Gemma, Phi-3 broken-template rows all repair FIX-OK)
- tests/test_load_correct_tokenizer_pr4426.py: 5/5
- tests/test_chat_template_followups.py: 18/18 (T11-T18 cover
non-ChatML repair + probe failure modes)
- tests/test_mistral_pr4426.py: 5/5 byte-identical
- tests/test_qwen_pr4426.py: 14/14 byte-identical (Qwen3-Guard AST
gate still rejects)
- tests/hermes3_lora_pr4426.py reload: patched template ends with
`<|im_start|>assistant\n`, inference returns sensible output.
- temp/sim/battery.py: 79/79 followup; vs baseline: 0 regressions,
9 improvements.
- Spot-check probe on real stripped tokenizers (Hermes-3, Phi-4,
Llama-3.2-1B, Gemma-3-1B): all derive the expected prefix.
for more information, see https://pre-commit.ci
…comment-safe end scan Resolves three reviewer findings on PR unslothai#5049 (`fix/chat-template-followups`): Finding #1 [10/10]: dict/list variants now route through `_fix_chat_template_for_tokenizer` via a new `_VariantTokenizerProxy` adapter. Previously the dict/list branches called `_fix_chat_template` directly, silently bypassing the warn/strict (`UNSLOTH_STRICT_CHAT_TEMPLATE`) contract, the `no == yes` diagnostic, broken-existing-block detection, and `_validate_patched_template` guard. The proxy swaps `base.chat_template` to the variant string before each `apply_chat_template` call so tokenizer globals (`bos_token`, custom filters, `raise_exception`) remain available; if the base is read-only it falls back to isolated Jinja rendering. Finding #2 [1/10]: `_has_add_generation_prompt_block` now requires the `If` body to contain at least one `Output` node (a new `_if_body_emits_content` helper walks descendants). This distinguishes a real generation-prompt block from a header guard like `{% if not add_generation_prompt is defined %}{% set ... %}{% endif %}` (body contains only `Assign`) which references the name but emits nothing. Also dropped a now-redundant `"add_generation_prompt" not in scrubbed` guard in `_fix_chat_template` Case 2 so header-guarded templates still get repaired. Finding #4 [1/10]: `_find_end_position` now replaces Jinja comments with equal-length whitespace before scanning for `{% endfor %}` / `{% endif %}` tokens. This prevents a trailing comment containing those tokens from being picked as the real end tag. Positions in the padded string map 1:1 to positions in the original template. Tests: - tests/test_chat_template_followups.py: 21/21 (T19 strict-mode dict variant, T20 header-guard repair, T21 comment-endfor trap added; T4/T5 stubs updated with a working apply_chat_template that routes through Jinja). - tests/test_fix_chat_template_pr4426.py: 51/51 cells unchanged. - tests/test_load_correct_tokenizer_pr4426.py: 5/5. - tests/test_mistral_pr4426.py: 5/5 byte-identical. - tests/test_qwen_pr4426.py: 14/14 byte-identical. - temp/sim/battery.py: 79/79 followup; 0 regressions vs baseline. - Phase 3 Hermes-3 broken-LoRA reload: inference still returns `'The answer to the equation 2+2 is 4.'`. - Spot-checks on Hermes-3 / Phi-4 / Llama-3.2-1B / Gemma-3-1B real stripped templates: probe still derives the expected prefix.
for more information, see https://pre-commit.ci
Pure comment minimization across `_find_end_position`, `_has_add_generation_prompt_block`, `_if_body_emits_content`, `_derive_assistant_prefix_by_render`, `_fix_chat_template` Case 2, and `_VariantTokenizerProxy`. No behavior change; same intent, fewer lines. All 21 follow-up tests and the 51-cell Phase 1 matrix still pass.
for more information, see https://pre-commit.ci
Three real bugs from the 10-agent Opus review:
1. Probe now uses `jinja2.sandbox.SandboxedEnvironment` instead of bare
`jinja2.Environment`. The probe renders at model-load time (before
the user calls `apply_chat_template`), so it was a new eager
code-execution surface that the base HF tokenizer loading does not
have. SandboxedEnvironment blocks attribute-chain exploits at
negligible cost.
2. `_repair_string_template` now tries validation with both
`is_sharegpt=False` and `is_sharegpt=True`. Previously, when
`_fix_chat_template` internally fell back to the other schema via
its dual-probe, the outer validation still used the caller's
original `is_sharegpt` -- rendering with the wrong message keys and
spuriously dropping a valid repair.
3. `_has_add_generation_prompt_block` now skips `If` nodes whose test
is a `Not` expression. A negated gate like
`{% if not add_generation_prompt %}{{ x }}{% endif %}` fires when
agp=False, so its emitting body is not a generation block -- but the
old code counted any Name reference regardless of polarity.
Cleanup: removed unused `self._label`, added `\r` escape in
generation-block literal, switched variant labels to `!r` formatting,
removed redundant `import os as _os`.
for more information, see https://pre-commit.ci
|
/gemini review |
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request refactors the chat template repair logic by introducing AST-based detection of generation prompt blocks and a rendering-based method for deriving assistant prefixes. It also adds support for multi-variant templates and enhances error handling. Reviewers suggested performance improvements for Jinja environment reuse and AST traversal, along with a readability refactor for format detection.
| try: | ||
| import jinja2 | ||
| import jinja2.nodes | ||
|
|
||
| ast = jinja2.Environment().parse(chat_template) |
There was a problem hiding this comment.
For a minor performance improvement, you could create a single jinja2.Environment() instance at the module level and reuse it across _template_ends_with_toplevel_for, _has_add_generation_prompt_block, _derive_assistant_prefix_by_render, and _VariantTokenizerProxy. This avoids repeated object creation.
For example, you could define _JINJA_ENV = jinja2.Environment() at the module level and then change this line to:
ast = _JINJA_ENV.parse(chat_template)| if any( | ||
| isinstance(d, jinja2.nodes.Output) | ||
| for d in node.find_all(jinja2.nodes.Output) | ||
| ): | ||
| return True |
There was a problem hiding this comment.
Using any with find_all traverses the entire subtree, which is slightly inefficient. For better performance, you can use node.find(...) which stops after the first match. The isinstance check is also redundant since find_all already filters by type.
if node.find(jinja2.nodes.Output) is not None:
return True| try: | ||
| messages = [ | ||
| {"role": "user", "content": "Who are you?"}, | ||
| ] | ||
| tokenizer.apply_chat_template( | ||
| messages, add_generation_prompt = False, tokenize = False | ||
| [{"role": "user", "content": "Who are you?"}], | ||
| add_generation_prompt = False, | ||
| tokenize = False, | ||
| ) | ||
| is_sharegpt = False | ||
| except: | ||
| except Exception: | ||
| try: | ||
| messages = [ | ||
| {"from": "human", "value": "Who are you?"}, | ||
| ] | ||
| tokenizer.apply_chat_template( | ||
| messages, add_generation_prompt = False, tokenize = False | ||
| [{"from": "human", "value": "Who are you?"}], | ||
| add_generation_prompt = False, | ||
| tokenize = False, | ||
| ) | ||
| is_sharegpt = True | ||
| except: | ||
| except Exception: | ||
| is_sharegpt = None |
There was a problem hiding this comment.
This nested try-except block for detecting the chat format can be refactored for better readability. Consider using a small helper function to probe for each format. For example:
def _check_format(messages):
try:
tokenizer.apply_chat_template(messages, add_generation_prompt=False, tokenize=False)
return True
except Exception:
return False
is_sharegpt = None
if _check_format([{"role": "user", "content": "Who are you?"}]):
is_sharegpt = False
elif _check_format([{"from": "human", "value": "Who are you?"}]):
is_sharegpt = True1. Import jinja2.sandbox explicitly -- bare `import jinja2` does not populate the sandbox submodule, so SandboxedEnvironment raised AttributeError (swallowed by except), silently disabling all Case-2 (GH#4150) chat-template repairs. 2. Distinguish "has agp block but broken" from "no agp block" in the user-facing warning. The has-block path was reusing the no-block message, telling users the block was missing when it was present but non-functional.
|
/gemini review |
There was a problem hiding this comment.
Code Review
This pull request significantly enhances the chat template repair logic in unsloth/tokenizer_utils.py. It introduces AST-based template analysis to ensure structural correctness, a robust 'render-diff' approach for automatically inferring assistant prefixes, and support for multi-variant templates (dictionaries or lists). Feedback includes a recommendation to simplify a redundant node type check in _if_body_emits_content and a suggestion to update the keyword argument handling in the _VariantTokenizerProxy fallback to use conversation for better compatibility with standard API expectations.
| if any( | ||
| isinstance(d, jinja2.nodes.Output) | ||
| for d in node.find_all(jinja2.nodes.Output) | ||
| ): | ||
| return True |
There was a problem hiding this comment.
The isinstance(d, jinja2.nodes.Output) check is redundant because node.find_all(jinja2.nodes.Output) already returns an iterator of jinja2.nodes.Output nodes. The any() call can be made directly on the result of find_all().
if any(node.find_all(jinja2.nodes.Output)):
return True| messages = args[0] if args else kwargs.get("messages", []) | ||
| add_generation_prompt = kwargs.get("add_generation_prompt", False) | ||
| return env.from_string(self._template).render( | ||
| messages = messages, | ||
| add_generation_prompt = add_generation_prompt, | ||
| ) |
There was a problem hiding this comment.
The keyword argument for apply_chat_template is conversation, not messages. While the current usage within this file is positional, using the correct keyword name makes this fallback logic robust for other potential uses.
| messages = args[0] if args else kwargs.get("messages", []) | |
| add_generation_prompt = kwargs.get("add_generation_prompt", False) | |
| return env.from_string(self._template).render( | |
| messages = messages, | |
| add_generation_prompt = add_generation_prompt, | |
| ) | |
| messages = args[0] if args else kwargs.get("conversation", []) | |
| add_generation_prompt = kwargs.get("add_generation_prompt", False) | |
| return env.from_string(self._template).render( | |
| messages = messages, | |
| add_generation_prompt = add_generation_prompt, | |
| ) |
1. Sandbox fallback Jinja env in _VariantTokenizerProxy.apply_chat_template
(use SandboxedEnvironment, matching _derive_assistant_prefix_by_render)
2. Unwrap benign outer-If guards in _template_ends_with_toplevel_for so
templates like {% if messages %}{% for ... %}{% endfor %}{% endif %}
are still repairable (preserves Qwen3-Guard rejection via else-branch
and add_generation_prompt-name checks)
3. Preserve raw name_or_path in _VariantTokenizerProxy._source_path so
local-path detection works for dict/list variant tokenizers
4. Context-aware strict-mode messages: omit "will still load" and
"Set UNSLOTH_STRICT_CHAT_TEMPLATE=1" when already raising
|
/gemini review |
There was a problem hiding this comment.
Code Review
This pull request significantly enhances the chat template repair logic in unsloth/tokenizer_utils.py by introducing AST-based structural analysis and a render-diff heuristic to automatically detect assistant prefixes. It also adds support for multi-variant templates via a new proxy class and improves user-facing warning messages. Feedback suggests making the fallback rendering logic more robust by supporting the conversation keyword and passing extra template variables, as well as using more flexible string matching when wrapping trailing expressions. Additionally, several staging artifacts and metadata files were identified for removal.
| messages = args[0] if args else kwargs.get("messages", []) | ||
| add_generation_prompt = kwargs.get("add_generation_prompt", False) | ||
| return env.from_string(self._template).render( | ||
| messages = messages, | ||
| add_generation_prompt = add_generation_prompt, | ||
| ) |
There was a problem hiding this comment.
The fallback rendering logic in _VariantTokenizerProxy.apply_chat_template has two issues:
- It uses
kwargs.get("messages"), but the standard keyword argument in the Transformersapply_chat_templateAPI isconversation. If a user calls it withconversation=msgs, the fallback will default to an empty list. - It only passes
messagesandadd_generation_promptto the Jinjarendermethod. Any extra keyword arguments passed toapply_chat_template(which are intended for the template, such asdate_stringin Llama 3.1) will be lost.
While this proxy is primarily used for internal repair logic with simple messages, making it more robust ensures that validation and potential future uses work correctly for all templates.
| messages = args[0] if args else kwargs.get("messages", []) | |
| add_generation_prompt = kwargs.get("add_generation_prompt", False) | |
| return env.from_string(self._template).render( | |
| messages = messages, | |
| add_generation_prompt = add_generation_prompt, | |
| ) | |
| messages = args[0] if args else kwargs.get("conversation", kwargs.get("messages", [])) | |
| add_generation_prompt = kwargs.get("add_generation_prompt", False) | |
| render_kwargs = {**kwargs, "messages": messages, "add_generation_prompt": add_generation_prompt} | |
| for key in ("tokenize", "padding", "truncation", "max_length", "return_tensors", "return_dict", "return_assistant_tokens_mask", "conversation"): | |
| render_kwargs.pop(key, None) | |
| return env.from_string(self._template).render(**render_kwargs) |
| and after_endfor.startswith("{{") | ||
| and after_endfor.endswith("}}") |
There was a problem hiding this comment.
The check for a trailing expression in Case 1 is very strict and will fail if there is any whitespace (like a newline) between the {% endfor %} tag and the {{ ... }} expression. Using .strip() or .lstrip()/.rstrip() would make this repair logic more robust for templates that are formatted with newlines.
| and after_endfor.startswith("{{") | |
| and after_endfor.endswith("}}") | |
| and after_endfor.lstrip().startswith("{{") | |
| and after_endfor.rstrip().endswith("}}") |
| @@ -0,0 +1,718 @@ | |||
| diff --git a/unsloth/tokenizer_utils.py b/unsloth/tokenizer_utils.py | |||
There was a problem hiding this comment.
This file appears to be a staging artifact containing a diff of the changes. It should be removed from the repository before merging to keep the codebase clean.
| { | ||
| "n_files_with_reverts": 0, | ||
| "n_total_reverted_lines": 0, | ||
| "severity": "none", | ||
| "reverts": [], | ||
| "owner_repo": "unslothai/unsloth", | ||
| "pr_number": 5049, | ||
| "base_ref": "main", | ||
| "pr_head_ref": "HEAD" | ||
| } No newline at end of file |
There was a problem hiding this comment.
This file seems to be metadata from a staging or CI process and likely does not belong in the final pull request.
|
/gemini review |
There was a problem hiding this comment.
Code Review
This pull request refactors and enhances the chat template repair mechanism in unsloth/tokenizer_utils.py. Key improvements include AST-based template parsing to identify structural nodes, support for multi-variant tokenizers through a new proxy class, and a render-diff approach to automatically derive assistant prefixes. The update also introduces a strict mode for template validation and updates several subproject dependencies. Feedback suggests that the fallback Jinja rendering path should include common tokens like bos_token and eos_token to prevent potential errors in templates that depend on them.
| return env.from_string(self._template).render( | ||
| messages = messages, | ||
| add_generation_prompt = add_generation_prompt, | ||
| ) |
There was a problem hiding this comment.
In the fallback path for read-only tokenizers, the Jinja rendering environment only receives messages and add_generation_prompt. Many real-world chat templates rely on other variables like bos_token, eos_token, or pad_token. Without these, the rendering will likely fail with an UndefinedError during the repair diagnostic. Consider passing these common tokens from self._base to the render call to improve the robustness of the fallback path.
return env.from_string(self._template).render(
messages = messages,
add_generation_prompt = add_generation_prompt,
bos_token = getattr(self._base, "bos_token", ""),
eos_token = getattr(self._base, "eos_token", ""),
pad_token = getattr(self._base, "pad_token", ""),
unk_token = getattr(self._base, "unk_token", ""),
)
…hai#5056) * Fix review findings for PR #49 1. Sandbox fallback Jinja env in _VariantTokenizerProxy.apply_chat_template (use SandboxedEnvironment, matching _derive_assistant_prefix_by_render) 2. Unwrap benign outer-If guards in _template_ends_with_toplevel_for so templates like {% if messages %}{% for ... %}{% endfor %}{% endif %} are still repairable (preserves Qwen3-Guard rejection via else-branch and add_generation_prompt-name checks) 3. Preserve raw name_or_path in _VariantTokenizerProxy._source_path so local-path detection works for dict/list variant tokenizers 4. Context-aware strict-mode messages: omit "will still load" and "Set UNSLOTH_STRICT_CHAT_TEMPLATE=1" when already raising * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
1 participant
Staging mirror of unslothai#5049
Original PR: unslothai#5049
Author: danielhanchen
This is a staging copy for review and editing. Once finalized, changes will be pushed back to the original PR.
Original description
Follow-up hardening on top of PR unslothai#4426 (which fixed the unslothai#4150 RuntimeError for ChatML LoRA reloads). Addresses the six follow-ups identified during the unslothai#4426 investigation.
Behavior changes
1. Warn-by-default instead of RuntimeError
Before: if the tokenizer's chat_template ignored
add_generation_promptand_fix_chat_templatecould not repair it,fix_chat_templateraised a hard RuntimeError that blocked model loading entirely. In the unslothai#4150 scenario this meant a tokenizer saved by a downstream tool (LlamaFactory, Axolotl) could not be reloaded at all.After: repair failure emits a
logger.warning_onceand returns the original template. Model loading continues; the user can either manually settokenizer.chat_templateor supply their own prompt for inference.UNSLOTH_STRICT_CHAT_TEMPLATE=1restores the pre-warn hard fail for users who want it.2. Local path vs HF hub distinguished in the message
Before:
Please file a bug report to the maintainers of <name_or_path>. In the unslothai#4150 case<name_or_path>was the user's ownsaves/Llama-3.1-8B-Instruct/lora/train_.../folder -- the user could not file a bug against their own directory.After: messages are branched on
os.path.isdir(name_or_path). Local paths get a message pointing at the likely downstream tool that re-serialized the tokenizer; HF hub IDs get a message pointing at the upstream model maintainers.3. Dict / list chat_template now handled
Before: passing a tokenizer whose
chat_templateis a dict (e.g. Hermes-3{default, tool_use}) or list to_fix_chat_templatecrashed withAttributeError: 'dict' object has no attribute 'find'. Surfaced during unslothai#4426 Phase 3 testing.After:
fix_chat_templatedetects dict / list shape and fixes each variant independently, preserving the outer structure.Internals
_find_end_positionnow matches all four Jinja whitespace-control variants ({% %},{%- %},{% -%},{%- -%}) and returns the rightmost endfor/endif. Previously{%- endfor -%}(both-side dash, used by Qwen3-Guard) was silently bypassed because the oldstr.findwas only looking for two specific dash styles._has_add_generation_prompt_blockuses Jinja AST (jinja2.nodes.If/Namewalk) instead of substring matching. Templates that mentionadd_generation_promptonly inside a comment are correctly classified as "no block"; templates that use whitespace-control variants on theif/endiftags are correctly classified as "has block"._template_ends_with_toplevel_forgates the GH#4150 ChatML repair on the AST: only fires when the last structural top-level node is a For loop (the standard ChatML shape), ignoring trailing pure-whitespace Output nodes. Templates wrapped in an outer If (Qwen3-Guard) are now explicitly skipped at the_fix_chat_templatelevel, not just atload_correct_tokenizer's name-based exemption. This means direct callers of_fix_chat_templatealso get correct behavior._validate_patched_template